"""
Copyright (c) 2022 Huawei Technologies Co.,Ltd.

openGauss is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:

          http://license.coscl.org.cn/MulanPSL2

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
"""
"""
Case Type   : 数据库系统
Case Name   : 打开快速负载均衡，同一集群，1个url连接2个节点、 1个url连接3个节点，连接database、username不同
Create At   : 2023/04/12
Owner       : @peilinqian
Description :
    step0: 检查集群状态，不正常时重启;
    step1: 创建指定数据库;
    step2: 创建指定用户;
    step3: stop其中一个节点;
    step4: 远程复制java脚本至指定主机;
    step5: 编译java文件;
    step6: 以线程运行java脚本;
    step7: 等待5s，查询两个正常节点连接数量;
    step8: 等待18s，恢复stop的节点;
    step9: java脚本执行70s左右时，查询两个节点连接数量;
    step10: java脚本执行82s左右新建连接后，查询节点连接数量;
    step11：验证step5的Java执行结果日志;
Expect      :
    step0: 检查集群状态，不正常时重启; expect:状态正常
    step1: 创建指定数据库; expect:创建成功
    step2: 创建指定用户; expect:创建成功
    step3: stop其中一个节点; expect:stop成功
    step4: 远程复制java脚本至指定主机; expect:复制成功
    step5: 编译java文件; expect:编译成功
    step6: 以线程运行java脚本; expect: 操作成功
    step7: 等待5s，查询两个正常节点连接数量; expect: 查询结果21/9
    step8: 等待18s，恢复stop的节点; expect: 恢复成功
    step9: java脚本执行70s左右时，查询两个节点连接数量; expect: 连接数量均为0
    step10: java脚本执行82s左右新建连接后，查询节点连接数量; expect: 连接数量为12/12/6
    step11：验证step5的Java执行结果日志; expect:包含节点探活及缓存线程、关闭连接线程、多次kill关键日志
History     :
    modified by @peilinqian at 2023-5-30: 建连时间不确定导致未进行快速负载，优化用例增加恢复数据库前的等待时间
"""
import os
import re
import time
import unittest

from testcase.utils.ComThread import ComThread
from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node
from yat.test import macro

pri_sh = CommonSH('PrimaryDbUser')


@unittest.skipIf(pri_sh.get_node_num() < 3, '非1主2备环境不执行')
class JDBCQuikAutoBalance(unittest.TestCase):

    def setUp(self):
        self.log = Logger()
        self.log.info("----this is setup----")
        self.log.info(f'----{os.path.basename(__file__)}:start----')
        self.pri_root = Node(node='PrimaryRoot')
        self.sta1_sh = CommonSH('Standby1DbUser')
        self.sta2_sh = CommonSH('Standby2DbUser')
        self.db_name1 = 'db_jdbc_quickautobalance_case0025_1'
        self.db_name2 = 'db_jdbc_quickautobalance_case0025_2'
        self.u_name1 = 'u_jdbc_quickautobalance_case0025_1'
        self.u_name2 = 'u_jdbc_quickautobalance_case0025_2'
        self.target_path = os.path.join(macro.DB_BACKUP_PATH, 'JDBC_test')
        self.java_name = "JDBC_QuikAutoBalance_Case0025"
        self.com = Common()
        self.constant = Constant()
        self.assert_flag1 = 'Start scheduleExecutorService'
        self.assert_flag2 = 'Start closeConnectionScheduledFuture'
        self.assert_flag3 = 'put ?.* idle connections into ' \
                            'abandonedConnectionList'

    def test_main(self):
        step_txt = '----step0: 检查集群状态，不正常时重启; expect:状态正常----'
        self.log.info(step_txt)
        if not pri_sh.get_db_cluster_status('status'):
            self.assertTrue(pri_sh.restart_db_cluster(), '执行失败' + step_txt)

        step_txt = '----step1: 创建指定数据库; expect:创建成功----'
        self.log.info(step_txt)
        create_sql = f"drop database if exists {self.db_name1};" \
            f"create database {self.db_name1};" \
            f"drop database if exists {self.db_name2};" \
            f"create database {self.db_name2};"
        create_result = pri_sh.execut_db_sql(create_sql)
        self.log.info(create_result)
        self.assertEqual(2, create_result.count(
            self.constant.CREATE_DATABASE_SUCCESS), '执行失败' + step_txt)

        step_txt = '----step2: 创建指定用户; expect:创建成功----'
        self.log.info(step_txt)
        create_sql = f"drop user if exists {self.u_name1};" \
            f"create user {self.u_name1} with " \
            f"password '{pri_sh.node.db_password}';" \
            f"grant all privileges to {self.u_name1};" \
            f"drop user if exists {self.u_name2};" \
            f"create user {self.u_name2} with " \
            f"password '{pri_sh.node.db_password}';" \
            f"grant all privileges to {self.u_name2}; checkpoint;"
        create_result = pri_sh.execut_db_sql(create_sql)
        self.log.info(create_result)
        self.assertEqual(2, create_result.count(
            self.constant.CREATE_ROLE_SUCCESS_MSG), '执行失败' + step_txt)
        time.sleep(2)

        step_txt = '----step3: stop其中一个节点; expect:stop成功----'
        self.log.info(step_txt)
        stop_result = pri_sh.stop_db_instance()
        self.assertIn(self.constant.GS_CTL_STOP_SUCCESS_MSG, stop_result,
                      '执行失败' + step_txt)

        step_txt = '----step4: 远程复制java脚本至指定主机; expect:复制成功----'
        self.log.info(step_txt)
        self.com.scp_file(self.pri_root, f'{self.java_name}.java',
                          self.target_path)
        check_cmd = f'ls {self.target_path}'
        check_result = self.pri_root.sh(check_cmd).result()
        self.assertIn(self.java_name, check_result, '执行失败' + step_txt)

        step_txt = '----step5: 编译java文件; expect:编译成功----'
        self.log.info(step_txt)
        self.pri_root.scp_put(macro.JDBC_PATH, os.path.join(self.target_path,
                                                            'postgresql.jar'))
        javac_cmd = f"javac -encoding utf-8 -cp " \
            f"{os.path.join(self.target_path, 'postgresql.jar')} " \
            f"{os.path.join(self.target_path, f'{self.java_name}.java')}"
        self.log.info(javac_cmd)
        javac_result = self.pri_root.sh(javac_cmd).result()
        self.log.info(javac_result)
        self.assertEqual('', javac_result, "执行失败" + step_txt)

        step_txt = '----step6: 以线程运行java脚本; expect: 操作成功----'
        self.log.info(step_txt)
        session1 = None
        try:
            url_host1 = f"{pri_sh.node.db_host}:{pri_sh.node.db_port}," \
                f"{self.sta1_sh.node.db_host}:{self.sta1_sh.node.db_port}"
            url_host2 = f"{pri_sh.node.db_host}:{pri_sh.node.db_port}," \
                f"{self.sta1_sh.node.db_host}:{self.sta1_sh.node.db_port}," \
                f"{self.sta2_sh.node.db_host}:{self.sta2_sh.node.db_port}"
            url_param = "loggerLevel=info\&autoBalance=leastconn" \
                        "\&enableQuickAutoBalance=true"
            java_cmd = f"java -cp " \
                f"{os.path.join(self.target_path, 'postgresql.jar')}" \
                f":{self.target_path} {self.java_name} {url_host1} " \
                f"{self.db_name1} {self.u_name1} {pri_sh.node.db_password} " \
                f"{url_host2} {self.db_name2} {self.u_name2} " \
                f"{pri_sh.node.db_password} {url_param} {url_param}"
            self.log.info(java_cmd)
            session1 = ComThread(self.com.get_sh_result,
                                 args=(self.pri_root, java_cmd,))
            session1.setDaemon(True)
            session1.start()
            start_time = time.time()

            step_txt = '----step7: 等待5s，查询两个正常节点连接数量; expect: 查询结果21/9----'
            self.log.info(step_txt)
            time.sleep(5)
            check_sql = f"select * from pg_stat_activity where " \
                f"application_name = 'PostgreSQL JDBC Driver' " \
                f"and datname in ('{self.db_name1}','{self.db_name2}');"
            self.log.info('----备1节点连接数量----')
            sta1_result = self.sta1_sh.execut_db_sql(check_sql)
            self.log.info(sta1_result)
            self.log.info('----备2节点连接数量----')
            sta2_result = self.sta2_sh.execut_db_sql(check_sql)
            self.log.info(sta2_result)
            self.assertIn('21 rows', sta1_result, "执行失败" + step_txt)
            self.assertIn('9 rows', sta2_result, "执行失败" + step_txt)

            step_txt = '----step8: 等待18s，恢复stop的节点; expect: 恢复成功----'
            self.log.info(step_txt)
            time.sleep(18)
            start_result = pri_sh.start_db_instance()
            self.assertIn(self.constant.RESTART_SUCCESS_MSG, start_result,
                          '执行失败' + step_txt)

            step_txt = '----step9: java脚本执行70s左右时，查询两个节点连接数量; ' \
                       'expect: 连接数量均为0----'
            self.log.info(step_txt)
            end_time = time.time()
            if end_time - start_time < 70:
                self.log.info('--不超过70s--')
                time.sleep(70 - end_time + start_time)
            else:
                self.log.info('--超过70s--')
            self.log.info('----备1节点连接数量----')
            sta1_result = self.sta1_sh.execut_db_sql(check_sql)
            self.log.info(sta1_result)
            self.log.info('----备2节点连接数量----')
            sta2_result = self.sta2_sh.execut_db_sql(check_sql)
            self.log.info(sta2_result)
            self.assertIn('0 rows', sta1_result, "执行失败" + step_txt)
            self.assertIn('0 rows', sta2_result, "执行失败" + step_txt)

            step_txt = '----step10: java脚本执行82s左右新建连接后，查询节点连接数量; ' \
                       'expect: 连接数量为12/12/6----'
            self.log.info(step_txt)
            end_time = time.time()
            if end_time - start_time < 82:
                self.log.info('--不超过82s--')
                time.sleep(82 - end_time + start_time)
            else:
                self.log.info('--超过82s--')
            self.log.info('----主节点连接数量----')
            pri_result = pri_sh.execut_db_sql(check_sql)
            self.log.info(pri_result)
            self.log.info('----备1节点连接数量----')
            sta1_result = self.sta1_sh.execut_db_sql(check_sql)
            self.log.info(sta1_result)
            self.log.info('----备2节点连接数量----')
            sta2_result = self.sta2_sh.execut_db_sql(check_sql)
            self.log.info(sta2_result)
            self.assertIn('12 rows', pri_result, "执行失败" + step_txt)
            self.assertIn('12 rows', sta1_result, "执行失败" + step_txt)
            self.assertIn('6 rows', sta2_result, "执行失败" + step_txt)
        finally:
            if session1:
                step_txt = '----step11：验证step5的Java执行结果日志; ' \
                           'expect:包含节点探活及缓存线程、关闭连接线程、多次kill关键日志----'
                self.log.info(step_txt)
                session1.join()
                session1_result = session1.get_result()
                self.assertIn(self.assert_flag1, session1_result,
                              "执行失败" + step_txt)
                self.assertIn(self.assert_flag2, session1_result,
                              "执行失败" + step_txt)
                self.assertEqual(len(re.findall(self.assert_flag3,
                                                session1_result)), 2,
                                 "执行失败" + step_txt)

    def tearDown(self):
        self.log.info('----this is teardown----')
        self.log.info('----清理环境----')
        ps_cmd = f"ps ux |grep {self.java_name}|grep -v grep|" \
            f"tr -s ' '|cut -d ' ' -f 2"
        self.log.info(ps_cmd)
        ps_result = self.pri_root.sh(ps_cmd).result()
        self.log.info(ps_result)
        if ps_result != '':
            self.log.info('----清理java进程----')
            kill_cmd = f"ps ux |grep {self.java_name}|grep -v grep|" \
                f"tr -s ' '|cut -d ' ' -f 2|xargs kill -9"
            self.log.info(kill_cmd)
            self.pri_root.sh(kill_cmd)

        step_txt1 = '----删除数据库; expect:删除成功----'
        self.log.info(step_txt1)
        drop_sql = f"drop database if exists {self.db_name1};" \
            f"drop database if exists {self.db_name2};"
        drop_result = pri_sh.execut_db_sql(drop_sql)
        self.log.info(drop_result)

        step_txt2 = '----删除指定用户; expect:删除成功----'
        self.log.info(step_txt2)
        drop_sql = f"drop user if exists {self.u_name1};" \
            f"drop user if exists {self.u_name2};"
        drop_user_result = pri_sh.execut_db_sql(drop_sql)
        self.log.info(drop_user_result)

        step_txt3 = '----清理编译脚本; expect:清理成功----'
        self.log.info(step_txt3)
        file_rm_cmd = f'rm -rf {self.target_path};' \
            f'if [ -d {self.target_path} ]; ' \
            f'then echo "exists"; else echo "not exists"; fi'
        self.log.info(file_rm_cmd)
        file_rm_result = self.pri_root.sh(file_rm_cmd).result()
        self.log.info(file_rm_result)

        self.log.info('----清理环境断言----')
        self.assertEqual(2, drop_result.count(
            self.constant.DROP_DATABASE_SUCCESS), '执行失败' + step_txt1)
        self.assertEqual(2, drop_user_result.count(
            self.constant.DROP_ROLE_SUCCESS_MSG), '执行失败' + step_txt2)
        self.assertEqual('not exists', file_rm_result, "执行失败" + step_txt3)

        self.log.info(f'----{os.path.basename(__file__)}:end----')
